package com.gitee.wsl.mathematics.function.randomForest

import kotlin.math.ln
import kotlin.random.Random

/**
 * @author lidapeng
 * @description 随机森林
 * @date 3:50 下午 2020/2/22
 */
class RandomForest {
    private val random: Random = Random.Default
    private var forest: Array<Tree?>
    var trustTh: Float = 0.1f //信任阈值
    var trustPunishment: Float = 0.1f //信任惩罚


    constructor(treeNub: Int) {
        if (treeNub > 0) {
            forest = arrayOfNulls(treeNub)
        } else {
            throw Exception("Number of trees must be greater than 0")
        }
    }

    val model: RfModel
        get() { //获取模型
            val rfModel: RfModel = RfModel()
            val nodeMap: MutableMap<Int, Node?> = HashMap()
            for (i in forest.indices) {
                val node = forest[i]!!.rootNode
                nodeMap[i] = node
            }
            rfModel.nodeMap = nodeMap
            return rfModel
        }

    @Throws(Exception::class)
    fun forest(`object`: Any): Int { //随机森林识别
        val map: MutableMap<Int, Float> = HashMap<Int, Float>()
        for (i in forest.indices) {
            val tree = forest[i]
            val treeWithTrust: TreeWithTrust = tree!!.judge(`object`)
            val type: Int = treeWithTrust.type
            //System.out.println(type);
            val trust: Float = treeWithTrust.trust
            if (map.containsKey(type)) {
                map[type] = map[type]!! + trust
            } else {
                map[type] = trust
            }
        }
        var type = 0
        var nub = 0f
        for ((key, myNub) in map) {
            //System.out.println("type==" + entry.getKey() + ",nub==" + myNub);
            if (myNub > nub) {
                type = key
                nub = myNub
            }
        }
        if (nub < (forest.size * trustTh)) {
            type = 0
        }
        return type
    }

    //rf初始化
    @Throws(Exception::class)
    fun init(dataTable: DataTable) {
        //一棵树属性的数量
        if (dataTable.size > 4) {
            val kNub = ((ln(dataTable.size.toDouble())).toInt() / ln(2.0f)).toInt()
            //int kNub = dataTable.getSize() / 2;
            // System.out.println("knNub==" + kNub);
            for (i in forest.indices) {
                val tree = Tree(getRandomData(dataTable, kNub), trustPunishment)
                forest[i] = tree
            }
        } else {
            throw Exception("Number of feature categories must be greater than 3")
        }
    }

    @Throws(Exception::class)
    fun study() { //学习
        for (i in forest.indices) {
            //System.out.println("开始学习==" + i + ",treeNub==" + forest.length);
            val tree = forest[i]
            tree!!.study()
        }
    }

    fun insert(`object`: Any) { //添加学习参数
        for (i in forest.indices) {
            val tree = forest[i]
            tree!!.dataTable!!.insert(`object`)
        }
    }

    //从总属性列表中随机挑选属性kNub个属性数量
    @Throws(Exception::class)
    private fun getRandomData(dataTable: DataTable, kNub: Int): DataTable {
        val attr: Set<String> = dataTable.keyType
        val myName: MutableSet<String> = HashSet()
        val key = dataTable.key //结果
        val list: MutableList<String> = ArrayList<String>()
        for (name in attr) { //加载主键
            if (name != key) {
                list.add(name)
            }
        }
        for (i in 0 until kNub) {
            val index: Int = random.nextInt(list.size)
            myName.add(list[index])
            list.removeAt(index)
        }
        myName.add(key!!)
        //System.out.println(myName);
        val data: DataTable = DataTable(myName)
        data.key = key
        return data
    }
}
